#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
测试ChromaDB数据结构

本测试文件专门测试ChromaDB的数据结构、连接状态和数据完整性，
确保向量数据库中的数据格式正确且可以正常访问。
"""

import asyncio
import pytest
import sys
import os
from pathlib import Path
from loguru import logger

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from src.database.chromadb_client import ChromaDBClient
from src.config import load_config


class TestChromaDB数据结构:
    """测试ChromaDB数据结构的测试类"""
    
    @pytest.fixture
    async def setup_chromadb_client(self):
        """设置ChromaDB客户端"""
        try:
            # 加载配置
            config = load_config()
            
            # 创建ChromaDB客户端
            chromadb_client = ChromaDBClient(config.get('chromadb', {}))
            await chromadb_client.initialize()
            
            yield chromadb_client
            
            # 清理
            await chromadb_client.shutdown()
            
        except Exception as e:
            logger.error(f"设置ChromaDB客户端失败: {e}")
            raise
    
    async def test_chromadb连接状态(self, setup_chromadb_client):
        """测试ChromaDB连接状态"""
        chromadb_client = setup_chromadb_client
        
        # 检查客户端是否正确初始化
        assert chromadb_client.client is not None, "ChromaDB客户端未正确初始化"
        
        # 检查连接状态
        try:
            # 尝试获取心跳
            heartbeat = chromadb_client.client.heartbeat()
            logger.info(f"ChromaDB心跳检测: {heartbeat}")
            
        except Exception as e:
            logger.error(f"ChromaDB连接测试失败: {e}")
            raise
    
    async def test_documents集合存在性(self, setup_chromadb_client):
        """测试documents集合是否存在"""
        chromadb_client = setup_chromadb_client
        
        try:
            # 获取所有集合
            collections = chromadb_client.client.list_collections()
            collection_names = [col.name for col in collections]
            
            logger.info(f"现有集合: {collection_names}")
            
            # 检查documents集合是否存在
            assert 'documents' in collection_names, "documents集合不存在"
            
            # 获取documents集合
            documents_collection = chromadb_client.get_collection('documents')
            assert documents_collection is not None, "无法获取documents集合"
            
        except Exception as e:
            logger.error(f"检查documents集合失败: {e}")
            raise
    
    async def test_documents集合数据结构(self, setup_chromadb_client):
        """测试documents集合的数据结构"""
        chromadb_client = setup_chromadb_client
        
        try:
            # 获取documents集合
            documents_collection = chromadb_client.get_collection('documents')
            
            # 获取集合统计信息
            count = documents_collection.count()
            logger.info(f"documents集合文档数量: {count}")
            
            if count > 0:
                # 获取前几个文档查看数据结构
                results = documents_collection.get(
                    limit=3,
                    include=['embeddings', 'metadatas', 'documents']
                )
                
                logger.info(f"获取到{len(results['ids'])}个文档样本")
                
                # 检查数据结构
                assert 'ids' in results, "结果中缺少ids字段"
                assert 'embeddings' in results, "结果中缺少embeddings字段"
                assert 'metadatas' in results, "结果中缺少metadatas字段"
                assert 'documents' in results, "结果中缺少documents字段"
                
                # 检查每个文档的结构
                for i, doc_id in enumerate(results['ids']):
                    logger.info(f"文档{i+1} ID: {doc_id}")
                    
                    if results['embeddings'] and i < len(results['embeddings']):
                        embedding = results['embeddings'][i]
                        logger.info(f"文档{i+1} 向量维度: {len(embedding) if embedding else 0}")
                    
                    if results['metadatas'] and i < len(results['metadatas']):
                        metadata = results['metadatas'][i]
                        logger.info(f"文档{i+1} 元数据: {metadata}")
                    
                    if results['documents'] and i < len(results['documents']):
                        document = results['documents'][i]
                        logger.info(f"文档{i+1} 内容长度: {len(document) if document else 0}")
            
            else:
                logger.warning("documents集合为空，无法检查数据结构")
                
        except Exception as e:
            logger.error(f"检查documents集合数据结构失败: {e}")
            raise
    
    async def test_向量数据完整性(self, setup_chromadb_client):
        """测试向量数据的完整性"""
        chromadb_client = setup_chromadb_client
        
        try:
            # 获取documents集合
            documents_collection = chromadb_client.get_collection('documents')
            
            # 获取所有文档的向量数据
            results = documents_collection.get(
                include=['embeddings', 'metadatas']
            )
            
            total_docs = len(results['ids'])
            valid_embeddings = 0
            invalid_embeddings = 0
            
            logger.info(f"开始检查{total_docs}个文档的向量完整性")
            
            for i, doc_id in enumerate(results['ids']):
                if results['embeddings'] and i < len(results['embeddings']):
                    embedding = results['embeddings'][i]
                    
                    if embedding and len(embedding) > 0:
                        valid_embeddings += 1
                        
                        # 检查向量是否包含有效数值
                        if all(isinstance(x, (int, float)) for x in embedding[:5]):  # 检查前5个元素
                            continue
                        else:
                            logger.warning(f"文档{doc_id}的向量包含无效数值")
                            invalid_embeddings += 1
                    else:
                        logger.warning(f"文档{doc_id}的向量为空")
                        invalid_embeddings += 1
                else:
                    logger.warning(f"文档{doc_id}缺少向量数据")
                    invalid_embeddings += 1
            
            logger.info(f"向量完整性检查结果:")
            logger.info(f"  总文档数: {total_docs}")
            logger.info(f"  有效向量: {valid_embeddings}")
            logger.info(f"  无效向量: {invalid_embeddings}")
            
            # 断言至少有一些有效向量
            assert valid_embeddings > 0, "没有找到有效的向量数据"
            
        except Exception as e:
            logger.error(f"检查向量数据完整性失败: {e}")
            raise
    
    async def test_搜索功能基础测试(self, setup_chromadb_client):
        """测试ChromaDB的基础搜索功能"""
        chromadb_client = setup_chromadb_client
        
        try:
            # 获取documents集合
            documents_collection = chromadb_client.get_collection('documents')
            
            # 检查集合是否为空
            count = documents_collection.count()
            if count == 0:
                logger.warning("documents集合为空，跳过搜索测试")
                return
            
            # 执行基础查询测试
            try:
                # 使用文本查询
                results = documents_collection.query(
                    query_texts=["Python编程"],
                    n_results=3
                )
                
                logger.info(f"文本查询返回{len(results['ids'][0])}个结果")
                
                # 检查结果结构
                assert 'ids' in results, "查询结果缺少ids字段"
                assert 'distances' in results, "查询结果缺少distances字段"
                
            except Exception as e:
                logger.error(f"文本查询测试失败: {e}")
                
            # 如果有向量数据，测试向量查询
            try:
                # 获取一个示例向量
                sample_results = documents_collection.get(
                    limit=1,
                    include=['embeddings']
                )
                
                if (sample_results['embeddings'] and 
                    len(sample_results['embeddings']) > 0 and 
                    sample_results['embeddings'][0]):
                    
                    sample_embedding = sample_results['embeddings'][0]
                    
                    # 使用向量查询
                    vector_results = documents_collection.query(
                        query_embeddings=[sample_embedding],
                        n_results=3
                    )
                    
                    logger.info(f"向量查询返回{len(vector_results['ids'][0])}个结果")
                    
            except Exception as e:
                logger.error(f"向量查询测试失败: {e}")
                
        except Exception as e:
            logger.error(f"搜索功能基础测试失败: {e}")
            raise


if __name__ == "__main__":
    # 配置日志
    logger.remove()
    logger.add(
        sys.stdout,
        format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>",
        level="INFO"
    )
    
    # 运行测试
    pytest.main([__file__, "-v", "-s"])